'''
    @project: Masdevallia
    @author: Austin Baker
    @created: 2/22/2014
    @license: http://opensource.org/licenses/mit-license.php

    The MIT License (MIT)

    Copyright (c) 2014 Austin Baker
    
    Permission is hereby granted, free of charge, to any person obtaining a copy
    of this software and associated documentation files (the "Software"), to deal
    in the Software without restriction, including without limitation the rights
    to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
    copies of the Software, and to permit persons to whom the Software is
    furnished to do so, subject to the following conditions:
    
    The above copyright notice and this permission notice shall be included in
    all copies or substantial portions of the Software.
    
    THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
    THE SOFTWARE.
    
    About:
    This is a collection of utility functions for use with Maya's DAG nodes.
'''

'''
============================================================
--->>   IMPORT MODULES
============================================================
'''
import maya.OpenMaya as OpenMaya
import masdevallia.utility.MScriptUtility as MScriptUtility

'''
============================================================
--->>   UTILITY METHODS
============================================================
'''
def createNode( inNodeName=None, inNodeType=None, inNodeParent=None ):
    '''
    Creates a node of the given type.
    
    @param: inNodeName. String. Name of the new node.
    @param inNodeType: String. Type of node to create.
    @param inNodeParent: String or MObject. Name of node's parent.
    @return: The new node.
    '''
    if not isinstance( inNodeType, str ):
        OpenMaya.MGlobal.displayError( 'createNode argument inNodeType must be a string.' )
        return
        
    if inNodeParent:
        if isinstance( inNodeParent, str ):
            if objExists( inNodeParent ):
                parentObj = getMObject( inNodeParent )
            
            else:
                OpenMaya.MGlobal.displayError( 'Object {0} doesn\t exist'.format( inNodeParent ) )
                return

        elif isinstance( inNodeParent, OpenMaya.MObject ):
            parentObj = inNodeParent
            
        else:
            OpenMaya.MGlobal.displayError( 'createNode argument inNodeParent must be a String or MObject.' )
            return
    
    mdgMod = OpenMaya.MDGModifier()
            
    if inNodeType != 'network':
        # MDagModifier can't make network nodes. So we special case those later.
        dagMod = OpenMaya.MDagModifier()
        
        if inNodeParent:    
            node = dagMod.createNode( inNodeType, parentObj )
            dagMod.doIt()
        else:
            node = dagMod.createNode( inNodeType )
            dagMod.doIt()
        
    else:
        # Special casing network nodes.
        node = mdgMod.createNode( inNodeType )
    
    if not inNodeName:
        inNodeName = inNodeType+'1'
    
    mdgMod.renameNode( node, inNodeName )
    mdgMod.doIt()
    
    return node
    
def getMObject( inObjName ):
    '''
    Converts a string name into a MObject.
    
    @param inObjName: String. Name of node to retrieve as MObject.
    @return: MObject.
    @return: None.
    '''
    selList = OpenMaya.MSelectionList()
    
    # Try to select the object.
    if inObjName[-1:] is not '*':
        try:
            selList.add( inObjName )
            obj = OpenMaya.MObject()
            selList.getDependNode( 0, obj )
            
            return obj 
        
        except:
            pass
    
    # Try to select the object with a wild card. This is to catch a case where
    # two or more object's have the same name in the scene and are in the default
    # namespace. MSelectionList.add() doesn't handle this case. This always returns
    # the object whose partial path name matches the passed in object name.
    try:
        selList.add( inObjName+'*' )
        arrayNames = []
        selList.getSelectionStrings( arrayNames )
        
        if len( arrayNames ) > 1:    
            name = '|'+inObjName
            idx = [ i for i,value in enumerate(arrayNames) if name == value ]
            
            if idx:
                obj = OpenMaya.MObject()
                selList.getDependNode( idx[0], obj )
                
                return obj                
    
    except:
        pass
    
    return None

def getDagPath( inObject ):
    '''
    Converts a string name into a DagPath.
    
    @param inObject: String or MObject. If String then name of node to retrieve as DagPath.
    @return: DagPath.
    @return: None.
    '''
    if isinstance( inObject, str ):
        name = inObject
        
    elif isinstance( inObject, OpenMaya.MObject ):
        name = OpenMaya.MFnDependencyNode( inObject ).name()
        
    else:
        OpenMaya.MGlobal.displayError( 'getDagPath() inObject flag must be a string or MObject. Got {0}.'.format( type( inObject ) ) )

    selList = OpenMaya.MSelectionList()
    
    # Try to select the object.
    if name[-1:] is not '*':
        try:            
            selList.add( name )
            dag = OpenMaya.MDagPath()
            component = OpenMaya.MObject()
            selList.getDagPath( 0, dag, component )
            
            return dag
        
        except:
            pass
    
    # Try to select the object with a wild card. This is to catch a case where
    # two or more object's have the same name in the scene and are in the default
    # namespace. MSelectionList.add() doesn't handle this case. This always returns
    # the object whose partial path name matches the passed in object name.
    try:
        selList.add( name+'*' )
        arrayNames = []
        selList.getSelectionStrings( arrayNames )
        
        if len( arrayNames ) > 1:
            fullname = '|'+name
            idx = [ i for i,value in enumerate(arrayNames) if fullname == value ]
            
            if idx:
                dag = OpenMaya.MDagPath()
                component = OpenMaya.MObject()
                selList.getDagPath( idx[0], dag, component )
                
                return dag
            
    except:
        pass
    
    return None

def objExists( inObjName ):
    '''
    Checks to see if the object exists in the scene.
    
    @param inObjName: String. Name of the object to search for.
    @return: Bool. True if found in the scene.
    '''
    mObj = getMObject( inObjName )
    
    if mObj is not None:
        return True
    
    else:
        return False

def rootOf( inItem ):
    '''
    API based version of the mel command 'rootOf'.
     
    @param inItem: MDagPath. Dag path of a node.
    @return: String. Fullpath name of the given node's root node.
    '''
    if isinstance( inItem, OpenMaya.MDagPath ):
        dagNode = OpenMaya.MFnDagNode( inItem )    
        parent = dagNode.parent( 0 )
        parentName = OpenMaya.MFnDependencyNode( parent ).name()
    
        if parentName == 'world':
            dagPath = OpenMaya.MDagPath()
            dagNode.getPath( dagPath )
            
            return dagPath.fullPathName()
        
        else:
            parentDag = getDagPath( parentName )
            
            return rootOf( parentDag )
        
    else:
        OpenMaya.MGlobal.displayInfo( 'inItem is not a MDagPath. {0}'.format( inItem ) )
    
def getPathname( inObj, inFull ):
    '''
    Gets the desired path name of the object.
    
    @param inObj: String or MObject. If string then name of object.
    @param inFullName: Bool. True for full path name.
    @return: String. Path name.
    '''
    if isinstance( inObj, str ):
        name = inObj
        
    elif isinstance( inObj, OpenMaya.MObject ):
        name = OpenMaya.MFnDependencyNode( inObj ).name()
        
    else:
        OpenMaya.MGlobal.displayError( 'getPathname() inObj flag must be a string or MObject. Got {0}.'.format( type( inObj ) ) )
    
    # if inObj is a MObject then the name will be unicode which getDagPath() doesn't
    # like. Convert it to a string.
    dag = getDagPath( str(name) )
    
    # World doesn't have a full path name and will cause an error if we try and
    # get one. So we need to check for that first.
    if name != 'world':
        if inFull:
            # These return unicode by default so we convert to strings.
            return str( dag.fullPathName() )
        
        else:
            return str( dag.partialPathName() )
        
    return name        
        
def getAttributeByName( inObj, inAttrName ):
    '''
    Gets a node's attribute as an MObject.
    
    @param inObj: MObject. Node to get attribute from.
    @param inAttrName: String. Name of attribute to get from node.
    @return: MObject.
    @return: None.
    '''
    depFn = OpenMaya.MFnDependencyNode( inObj )
    
    try:
        attr = depFn.attribute( inAttrName )
        
        return attr
    
    except:
        return None

def createPlug( target_obj, plug_fullname, plug_shortname, plug_type ):
    '''
    Creates a plug.
    
    @param target_obj: MObject. Object to add the plug to.
    @param plug_fullname: String. Long name of the plug.
    @param plug_shortname: String. Short name of the plug.
    @param plug_type: Type of plug to create.
    '''    
    mdg_mod = OpenMaya.MDGModifier()
    
    # Create the attribute.
    attr = None    
    if plug_type == 'string':
        typed_attr = OpenMaya.MFnTypedAttribute()
        attr = typed_attr.create( plug_fullname, plug_shortname, OpenMaya.MFnData.kString )
    
    elif plug_type == 'int':
        num_attr = OpenMaya.MFnNumericAttribute()
        
    elif plug_type == 'matrix':
        matrix_attr = OpenMaya.MFnMatrixAttribute()
    
    # Add it to the object.
    mdg_mod.addAttribute( target_obj, attr )
    mdg_mod.doIt()
    
def getPlugByName( inObj, inPlugName ):
    '''
    Gets a node's plug as an MPlug.
    
    @param inObj: MObject. Node to get plug from.
    @param inPlugName: String. Name of plug to get from node.
    @return: MPlug.
    @return: None.
    '''
    depFn = OpenMaya.MFnDependencyNode( inObj )
    
    try:
        plug = depFn.findPlug( inPlugName )
        
        return plug
    
    except:
        return None

def connectNodes( source_obj, source_plug_name, destination_obj, destination_plug_name ):
    '''
    @param source_obj: MObject. Source object.
    @param source_plug_name: String. Name of plug on parent node.
    @param destination_obj: MObject. Destination object.
    @param destination_plug_name: String. Name of plug on child node.
    '''
    source_plug = getPlugByName( source_obj, source_plug_name )
    destination_plug = getPlugByName( destination_obj, destination_plug_name )
    MDGMod = OpenMaya.MDGModifier()
    MDGMod.connect( source_plug, destination_plug )
    MDGMod.doIt()
        
def getPlugConnections( plug, dest, src ):
    '''
    @param plug: MPlug.
    @param dest: Bool. True to check plug as a destination.
    @param src: Bool. True to check the plug as a source.
    @return: MPlugArray.
    '''
    plug_array = OpenMaya.MPlugArray()
    
    if plug.connectedTo( plug_array, dest, src ):
        return plug_array
    
    else:
        return None

def getPlugValue( inPlug ):
    '''
    Gets the value of the given plug.
    
    @param inPlug: MPlug. The node plug.
    @return: The value of the passed in node plug.
    '''
    pAttribute = inPlug.attribute()
    apiType = pAttribute.apiType()
    
    # Float Groups - rotate, translate, scale; Compounds
    if apiType in [ OpenMaya.MFn.kAttribute3Double, OpenMaya.MFn.kAttribute3Float, OpenMaya.MFn.kCompoundAttribute ]:
        result = []
        
        if inPlug.isCompound():
            for c in xrange( inPlug.numChildren() ):
                result.append( getPlugValue( inPlug.child( c ) ) )
                
            return result
    
    # Distance
    elif apiType in [ OpenMaya.MFn.kDoubleLinearAttribute, OpenMaya.MFn.kFloatLinearAttribute ]:
        return inPlug.asMDistance().asCentimeters()
    
    # Angle
    elif apiType in [ OpenMaya.MFn.kDoubleAngleAttribute, OpenMaya.MFn.kFloatAngleAttribute ]:
        return inPlug.asMAngle().asDegrees()
    
    # TYPED
    elif apiType == OpenMaya.MFn.kTypedAttribute:
        pType = OpenMaya.MFnTypedAttribute( pAttribute ).attrType()
        
        # Matrix
        if pType == OpenMaya.MFnData.kMatrix:   
            return OpenMaya.MFnMatrixData( inPlug.asMObject() ).matrix()
        
        # String
        elif pType == OpenMaya.MFnData.kString: 
            return inPlug.asString()
        
    # MATRIX
    elif apiType == OpenMaya.MFn.kMatrixAttribute:
        return OpenMaya.MFnMatrixData( inPlug.asMObject() ).matrix()
    
    # NUMBERS
    elif apiType == OpenMaya.MFn.kNumericAttribute:
        pType = OpenMaya.MFnNumericAttribute( pAttribute ).unitType()

        if pType == OpenMaya.MFnNumericData.kBoolean:
            return inPlug.asBool()
        
        elif pType in [ OpenMaya.MFnNumericData.kShort, OpenMaya.MFnNumericData.kInt, OpenMaya.MFnNumericData.kLong, OpenMaya.MFnNumericData.kByte ]:
            return inPlug.asInt()
        
        elif pType in [ OpenMaya.MFnNumericData.kFloat, OpenMaya.MFnNumericData.kDouble, OpenMaya.MFnNumericData.kAddr ]:
            return inPlug.asDouble()
        
    # Enum
    elif apiType == OpenMaya.MFn.kEnumAttribute:
        return inPlug.asInt()

def setPlugValue( inPlug, inValue ):
    '''
    Sets the given plug's value to the passed in value.

    @param inPlug: _MPlug_. The node plug.
    @param inValue: _Type_. Any value of any data type.
    '''
    plugAttribute = inPlug.attribute()
    apiType = plugAttribute.apiType()
    #print plugAttribute.apiTypeStr()
    
    # Float Groups - rotate, translate, scale
    if apiType in [ OpenMaya.MFn.kAttribute3Double, OpenMaya.MFn.kAttribute3Float ]:
        result = []
        
        if inPlug.isCompound():
            if isinstance( inValue, list ):
                for c in xrange( inPlug.numChildren() ):
                    result.append( setPlugValue( inPlug.child( c ), inValue[ c ] ) )
                    
                return result
            
            elif type( inValue ) == OpenMaya.MEulerRotation:
                setPlugValue( inPlug.child( 0 ), inValue.x )
                setPlugValue( inPlug.child( 1 ), inValue.y )
                setPlugValue( inPlug.child( 2 ), inValue.z )
                
            else:
                OpenMaya.MGlobal.displayError( '{0} :: Passed in value ( {1} ) is {2}. Needs to be type list.'.format( inPlug.info(), inValue, type( inValue ) ) )
    
    # Distance
    elif apiType in [ OpenMaya.MFn.kDoubleLinearAttribute, OpenMaya.MFn.kFloatLinearAttribute ]:
        if isinstance( inValue, float ):
            value =  OpenMaya.MDistance( inValue, OpenMaya.MDistance.kCentimeters )
            inPlug.setMDistance( value )
            
        else:
            OpenMaya.MGlobal.displayError( '{0} :: Passed in value ( {1} ) is {2}. Needs to be type float.'.format( inPlug.info(), inValue, type( inValue ) ) )
    
    # Angle
    elif apiType in [ OpenMaya.MFn.kDoubleAngleAttribute, OpenMaya.MFn.kFloatAngleAttribute ]:
        if isinstance( inValue, float ):
            value = OpenMaya.MAngle( inValue, OpenMaya.MAngle.kDegrees )
            inPlug.setMAngle( value )
            
        else:
            OpenMaya.MGlobal.displayError( '{0} :: Passed in value ( {1} ) is {2}. Needs to be type float.'.format( inPlug.info(), inValue, type( inValue ) ) )
            
    # Typed - matrix WE DON'T HANDLE THIS CASE YET!!!!!!!!!
    elif apiType == OpenMaya.MFn.kTypedAttribute:
        pType = OpenMaya.MFnTypedAttribute( plugAttribute ).attrType()
        
        if pType == OpenMaya.MFnData.kMatrix:
            if isinstance( inValue, OpenMaya.MPlug ):
                pass
            
            else:
                plugNode = inPlug.node()
                MFnTrans = OpenMaya.MFnTransform( plugNode )
                sourceMatrix = OpenMaya.MTransformationMatrix( inValue )#.asMatrix()
                MFnTrans.set( sourceMatrix )
        
        # String
        elif pType == OpenMaya.MFnData.kString:
            value = inValue
            inPlug.setString( value )
    
    # MATRIX
    elif apiType == OpenMaya.MFn.kMatrixAttribute:
        if isinstance( inValue, OpenMaya.MPlug ):
            # inValue must be a MPlug!
            sourceValueAsMObject = OpenMaya.MFnMatrixData( inValue.asMObject() ).object()
            inPlug.setMObject( sourceValueAsMObject )
            
        else:
            OpenMaya.MGlobal.displayError( 'Value object is not an MPlug. To set a MMatrix value, both passed in variables must be MPlugs.' )
    
    # Numbers
    elif apiType == OpenMaya.MFn.kNumericAttribute:
        pType = OpenMaya.MFnNumericAttribute( plugAttribute ).unitType()
        
        if pType == OpenMaya.MFnNumericData.kBoolean:
            if isinstance( inValue, bool ):
                inPlug.setBool( inValue )
                
            else:
                OpenMaya.MGlobal.displayError( '{0} :: Passed in value ( {1} ) is {2}. Needs to be type bool.'.format( inPlug.info(), inValue, type( inValue ) ) )
                
        elif pType in [ OpenMaya.MFnNumericData.kShort, OpenMaya.MFnNumericData.kInt, OpenMaya.MFnNumericData.kLong, OpenMaya.MFnNumericData.kByte ]:
            if isinstance( inValue, int ):
                inPlug.setInt( inValue )
            
            else:
                OpenMaya.MGlobal.displayError( '{0} :: Passed in value ( {1} ) is {2}. Needs to be type int.'.format( inPlug.info(), inValue, type( inValue ) ) )
        
        elif pType in [ OpenMaya.MFnNumericData.kFloat, OpenMaya.MFnNumericData.kDouble, OpenMaya.MFnNumericData.kAddr ]:
            if isinstance( inValue, float ):
                inPlug.setDouble( inValue )
           
            else:
                OpenMaya.MGlobal.displayError( '{0} :: Passed in value ( {1} ) is {2}. Needs to be type float.'.format( inPlug.info(), inValue, type( inValue ) ) )
                
    # Enums
    elif apiType == OpenMaya.MFn.kEnumAttribute:
        inPlug.setInt( inValue )

def getObjMatrix( inObj, inMatrixName ):
    '''
    Gets the desired matrix for the node.
    
    @param inObj: String or MObject. Node to get the matrix from.
    @param inMatrixName: String. Name of the matrix to get. Valid names are
        matrix, inverseMatrix, parentInverseMatrix, parentMatrix,
        worldInverseMatrix, worldMatrix, xformMatrix.
    @return: MMatrix.
    @return: None.
    '''
    if isinstance( inObj, str ):
        obj = getMObject( inObj )
        
    elif isinstance( inObj, OpenMaya.MObject ):
        obj = inObj
        
    else:
        OpenMaya.MGlobal.displayError( 'Object is type {0}. Needs to be string or MObject.'.format( type( obj ) ) )
        return None
    
    plug = getPlugByName( obj, inMatrixName )
    if plug is not None:
        if plug.isArray():
            plug = plug.elementByLogicalIndex( 0 )
            
        matrixData = OpenMaya.MFnMatrixData( plug.asMObject() )
        return OpenMaya.MMatrix( matrixData.matrix() )
    
    else:
        return plug

def doReparent( inChild, inParent, keepTransform=False ):
    '''
    Adds the given object(s) to this Node as children.
    
    @param inChildren: String, MObject or List (of string or MObject). Children to add to this Node.
    @param inParent: String or MObject. New parent for the children. There is a special case if 
            this value is World. Using this value will put inChild into the scene root.
    @param keepTransform: Bool. Default is False. If True the child object will keep it's transforms
            in world space. The default behavior of MDagModifier.reparentNode() is to keep
            the child objects local transforms the same after reparenting it to a new node.
            This results in the object shifting in world space.
    '''
    # This is a circular import. I should rewrite this function to not reference Node().
    import masdevallia.nodes.Node as Node
    
    newParent = None
    newChild = None
    
    if isinstance( inParent, str ):
        if inParent is 'World':
            newParent = OpenMaya.MObject().kNullObj
            
        else:
            newParent = getMObject( inParent )
            
    elif isinstance( inParent, OpenMaya.MObject ):
        newParent = inParent
        
    else:
        raise RuntimeError( 'inParent is not a string or MObject. {0}'.format( type( inParent ) ) )
    
    # Handle List case.
    if isinstance( inChild, list ):
        for child in inChild:
            if isinstance( child, OpenMaya.MObject ):
                newChild = child
                
            elif isinstance( child, str ):
                newChild = getMObject( child )    
                           
            else:
                raise RuntimeError( 'Child object is not a string or MObject. {0}'.format( type( child ) ) )
            
            # Store the pre-parenting world matrix of the child.
            childNode = Node.Node( newChild )
            origMatrix = childNode.getMatrixData( inMatrixType='worldMatrix' )
            
            # Do the parenting.
            dagMod = OpenMaya.MDagModifier()
            dagMod.reparentNode( newChild, newParent )
            dagMod.doIt()
            
            # Re-apply the original world matrix. 
            parentedChildNode = Node.Node( newChild )
            newMatrix = parentedChildNode.getMatrixData( inMatrixType='worldMatrix' )
            newMatrix.setTranslation( origMatrix.getTranslation() )
            newMatrix.setRotation( origMatrix.getRotation( OpenMaya.MQuaternion ) )
            
    # Handle single child cases.
    else:        
        if isinstance( inChild, str ):    
            newChild = getMObject( inChild )
            
        elif isinstance( inChild, OpenMaya.MObject ):
            newChild = inChild
            
        else:
            raise RuntimeError( 'inChild is not a list or string. {0}'.format( type( inChild ) ) )
        
        # Store the pre-parenting world matrix of the child.
        childNode = Node.Node( newChild )
        origMatrix = childNode.getMatrixData( inMatrixType='worldMatrix' )
        
        # Do the parenting.
        dagMod = OpenMaya.MDagModifier()
        dagMod.reparentNode( newChild, newParent )
        dagMod.doIt()
        
        # Re-apply the original world matrix. This will error here if the object
        # we are attempting to move is a shape.
        parentedChildNode = Node.Node( newChild )
        newMatrix = parentedChildNode.getMatrixData( inMatrixType='worldMatrix' )
        newMatrix.setTranslation( origMatrix.getTranslation() )
        newMatrix.setRotation( origMatrix.getRotation( OpenMaya.MQuaternion ) )
        
def getChildren( inObj, inType=None ):
    '''
    Returns a list of all children objects of the passed in object. The search can
    be limited to a single object type. Otherwise the search is for all children.
    
    @param inObj: MObject. Object to search for children.
    @param inType: MFn::Type. Type of child object to search for.
    @return List of Strings. Full path name of each child object.
    '''
    dagPath = getDagPath( inObj )
    dagFn = OpenMaya.MFnDagNode( dagPath )
    
    result = []
    for idx in xrange( dagPath.childCount() ):
        child_obj = dagFn.child( idx )
        if child_obj.apiType() != inType:
            child_dag_node = OpenMaya.MFnDagNode( child_obj )
            child_name = child_dag_node.fullPathName()
            result.extend( getChildren( child_dag_node.object(), inType ) )
        else:
            child_dag_node = OpenMaya.MFnDagNode( child_obj )
            child_name = child_dag_node.fullPathName()
            result.extend( getChildren( child_dag_node.object(), inType ) )
            result.append( child_name )
        
    return result

def getShapes( inObjectName ):
    '''
    Returns a list of all shapes under this kTransform.
    
    @param inObjectName: String. Name of object.
    @return: List of MObjects.
    '''
    results = []
    dag = getDagPath( inObjectName )
    
    # Get the number of shapes for this kTransform.
    numShapesPtr = MScriptUtility.getIntPointer()
    dag.numberOfShapesDirectlyBelow( numShapesPtr[0] )
    numShapes = numShapesPtr[1].getUint( numShapesPtr[0] )
    
    # Loop through the shapes.
    for shape_index in xrange( numShapes ):
        # We need to get the dag path every time so that we're
        # starting with the kTransform. Otherwise extendToShapeDirectlyBelow()
        # will only ever return the first shape regardless of the index
        # passed in. This is because it alters where the dag path points.
        temp_dag = getDagPath( inObjectName )
        temp_dag.extendToShapeDirectlyBelow( shape_index )
        results.append( temp_dag.node() )
                
    return results